Feature: NumPy-compliant distributed advanced indexing - #938

Description

TL;DR

This PR replaces Heat's legacy, local-only indexing with 100% NumPy-API-compliant advanced indexing capabilities across distributed nodes.

You can now seamlessly use boolean masks, integer arrays, negative-step slices etc. on distributed arrays without manual data shuffling:

import heat as ht

# Load a massive distributed dataset (e.g., 100 million rows, split across nodes/GPUs)
data = ht.random.randn(100_000_000, 50, split=0) 

#  distributed filtering condition 
outliers_mask = (data > 3.0).any(axis=1)

# extract only the outliers without gathering to a single node
outliers = data[outliers_mask]

This pull request introduces a significant overhaul of distributed indexing within dndarray.py, specifically targeting the __getitem__ and __setitem__ methods.

The logic has been completely refactored to identify zero-communication paths ("early out") for standard slices, while routing heavy, unordered (non-sequential) advanced indexing through highly optimized MPI collective communication.

Also, indexing.nonzero(), the kwarg as_tuple has been introduced (default: True) to comply with the Numpy API while giving users the choice to switch to torch-style output (2-D array). Merged separately in #2332

Main changes (LATEST UPDATE 18.9.2026)

dndarray.py

• replaced scattered key parsing with a centralized _resolve_indexing_state helper. This function torch-proofs all key inputs, handles broadcasting, aligns array dimension to indexed shape, and determines the indexing operation type for later dispatching. Returns a structured ProcessedKey NamedTuple.
• refactored the monolithic __getitem__ and __setitem__ functions. They are now wrappers that call the resolution state and dispatch to dedicated methods (e.g., __getitem_scalar, __setitem_mask, __getitem_advanced_local).
• added full support for non-sequential distributed advanced indexing (using integer arrays or boolean masks). It uses MPI.Alltoallv for cross-rank data fetching and assignment (__getitem_unordered and __setitem_unordered).
• introduced _resolve_duplicate_indices to guarantee NumPy-compliant "last assignment wins" semantics when using advanced indexing with duplicate indices on GPUs (thanks @Hakdag97 ).
• implemented support for descending/negative-step slices across distributed memory.
• added a __broadcast_value helper to automatically broadcast assigned values to match the target slice or boolean-mask shape during __setitem__ operations.
• updated __torch_proxy__ to explicitly track the split axis natively within the tensor's named dimensions for safer split axis tracking during dimensions-changing operations. discarded as PyTorch named tensors are no longer supported
• NEW updated __torch_proxy__ to use PyTorch meta tensors, to perform lightweight shape and index validation without allocating tensor memory
• refactored tests, removed orphaned legacy code and tests.
• introduced INDEXING.md in doc/source/ and added it to the .rst index this will be addressed in a different PR.

indexing.py

Changes to the indexing module have been merged with #2332.

Summary of distribution semantics (UPDATED 16.9.2026)

Note: Extracting a single element along the split axis will collapse that dimension, resulting in split=None.

Internal getitem/setitem routing logic

UPDATE 16.9.2026

graph TD
    Start((Receive Key)) --> CheckDist{Is array distributed?}
    
    CheckDist -- No --> LocalFastPath[Unwrap key & index underlying tensor directly]
    CheckDist -- Yes --> CheckScalar{Is key a pure scalar<br/>and not boolean?}
    
    CheckScalar -- Yes --> EvalRoot{Compute root rank}
    EvalRoot --> OpScalar[op_type = 'scalar']
    
    CheckScalar -- No --> CheckDistrMaskEarly{Is key a boolean mask<br/>aligned with array?}
    CheckDistrMaskEarly -- Yes --> OpDistrMask1[op_type = 'distr_mask']
    
    CheckDistrMaskEarly -- No --> ResolveKeys[Resolve key & check bounds]
    
    ResolveKeys --> AssessOpType{_assess_op_type}
    
    AssessOpType -->|root is not None| OpScalar[op_type = 'scalar']
    AssessOpType -->|split_key_is_ordered == 0| OpDist[op_type = 'distributed']
    AssessOpType -->|split_key_is_ordered == -1| OpDesc[op_type = 'descending_slice']
    AssessOpType -->|distr_mask_fast_path| OpDistrMask2[op_type = 'distr_mask']
    AssessOpType -->|key_is_mask_like| OpLocalMask[op_type = 'local_mask']
    AssessOpType -->|Default / Ordered / Slices| OpLocal[op_type = 'local']

    %% Map to actual handlers
    subgraph Handlers [Target dispatch methods]
        OpScalar --> H_Scalar[__getitem_scalar<br/>__setitem_scalar]
        OpDist --> H_Dist[__getitem_advanced_distributed<br/>__setitem_advanced_distributed]
        OpDesc --> H_Desc[__getitem_descending_slice_distributed<br/>__setitem_descending_slice_distributed]
        OpDistrMask1 & OpDistrMask2 --> H_DistMask[__getitem_mask<br/>__setitem_mask]
        OpLocalMask & OpLocal --> H_Local[__getitem_local<br/>__setitem_local]
    end
    
    %% Styling
    classDef target fill:#d4edda,stroke:#28a745,stroke-width:2px;
    class H_Scalar,H_Dist,H_Desc,H_DistMask,H_Local target;

Memory footprint

Scaling behaviour

Issue/s resolved: #703 #914 #918 #1012 #1019 #2135 #1816 #824

Type of change

• Breaking change

Memory requirements

Performance

Will follow

Due Diligence

• All split configurations tested
• Multiple dtypes tested in relevant functions
• Documentation updated (if needed)
• Updated changelog.md under the title "Pending Additions"

Does this change modify the behaviour of other functions? If so, which?

yes, everything that relied on the legacy indexing quirks (fixed) and everything that relied on 2D output from nonzero() (also fixed)

AI usage

We used frontier models (GPT, Gemini 3.1Pro - 3.8 extended), to help refactor, test, debug, optimize this PR, and to write the documentation.